/*
    This program is a research of implementing encrypting and decrypting messages using Elliptic Curve Cryptography (ECC) using Rust.
    Author: Smukx.E [@5mukx]
*/

// #![allow(deprecated)]
pub use k256::{elliptic_curve::{sec1::FromEncodedPoint, AffinePoint, Field}, EncodedPoint, ProjectivePoint, Scalar, Secp256k1};
pub use sha2::{Digest, Sha256};
pub use rand::rngs::OsRng;
pub use k256::elliptic_curve::group::GroupEncoding;
pub use k256::ecdsa::VerifyingKey;

fn encode_shellcode(
    shellcode: &[u8],
    public_key: &AffinePoint<Secp256k1>,
) -> (EncodedPoint, Vec<u8>) {
    let mut rng = OsRng;

    // generate the ephemeral keypair
    let k = Scalar::random(&mut rng);
    let r = (ProjectivePoint::generator() * k).to_affine();

    // compute shared secret
    let shared_secret = *public_key * k;
    let shared_secret_bytes = shared_secret.to_bytes();

    // derive encryption key from shared secret
    let mut hasher = Sha256::new();
    hasher.update(shared_secret_bytes);
    let encryption_key = hasher.finalize();

    // Encrypt shellcode
    let encrypted_shellcode: Vec<u8> = shellcode
        .iter()
        .zip(encryption_key.iter().cycle())
        .map(|(&byte, &key)| byte ^ key)
        .collect();

    (EncodedPoint::from(&r), encrypted_shellcode)
}

fn decode_shellcode(
    encrypted_shellcode: &[u8],
    r: &EncodedPoint,
    private_key: &Scalar,
) -> Vec<u8> {
    // Compute shared secret
    let r_point = ProjectivePoint::from_encoded_point(r).expect("Invalid R point");
    let shared_secret = r_point * private_key;
    let shared_secret_bytes = shared_secret.to_bytes();

    // derive decryption key from shared secret
    let mut hasher = Sha256::new();
    hasher.update(shared_secret_bytes);
    let decryption_key = hasher.finalize();

    // Decrypt shellcode
    encrypted_shellcode
        .iter()
        .zip(decryption_key.iter().cycle())
        .map(|(&byte, &key)| byte ^ key)
        .collect()
}

fn main() {
    // Example shellcode (placeholder)
    let shellcode: &[u8] = b"Hello, World!";

    // Generate ECC key pair
    let private_key = Scalar::random(&mut OsRng);
    let public_key = (ProjectivePoint::generator() * private_key).to_affine();

    println!("Private Key: {:?}", private_key);
    println!("Public Key: {:?}", public_key);

    // Convert AffinePoint to VerifyingKey (or PublicKey)
    VerifyingKey::from_encoded_point(&EncodedPoint::from(public_key))
        .expect("Invalid public key");

    let (r, encrypted_shellcode) = encode_shellcode(shellcode, &public_key);
    
    println!("Encrypted Shellcode: {:?}", encrypted_shellcode);

    // Decode the shellcode
    let decrypted_shellcode = decode_shellcode(&encrypted_shellcode, &r, &private_key);

    println!(
        "Decrypted Shellcode: {:?}",
        String::from_utf8(decrypted_shellcode).unwrap()
    );
}